from fastapi import APIRouter, HTTPException, Response
from pydantic import BaseModel, Field, Extra
import uuid
from server.middleware.log import logger
import os
import requests
import tempfile
import aiofiles

os.environ['OMP_NUM_THREADS'] = '1'

import signal
import sys
import warnings
import onnxruntime
from argparse import ArgumentParser, HelpFormatter

import facefusion.choices
import facefusion.globals
from facefusion.core import apply_args, pre_check, conditional_process
from facefusion import face_analyser, face_masker, content_analyser, config, metadata, logger, wording
from facefusion.processors.frame.core import get_frame_processors_modules, load_frame_processor_module
from facefusion.common_helper import create_metavar
from facefusion.execution_helper import encode_execution_providers
from facefusion.memory import limit_system_memory
from facefusion.filesystem import list_directory, clear_temp

onnxruntime.set_default_logger_severity(3)
warnings.filterwarnings('ignore', category = UserWarning, module = 'gradio')
warnings.filterwarnings('ignore', category = UserWarning, module = 'torchvision')


def destroy() -> None:
	if facefusion.globals.target_path:
		clear_temp(facefusion.globals.target_path)
	sys.exit(0)


def run_model():
	signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
	program = ArgumentParser(formatter_class=lambda prog: HelpFormatter(prog, max_help_position=130), add_help=False)
	# general
	program.add_argument('-s', '--source', help=wording.get('help.source'), action='append', dest='source_paths', default=config.get_str_list('general.source_paths'))
	program.add_argument('-t', '--target', help=wording.get('help.target'), dest='target_path', default=config.get_str_value('general.target_path'))
	program.add_argument('-o', '--output', help=wording.get('help.output'), dest='output_path', default=config.get_str_value('general.output_path'))
	program.add_argument('-v', '--version', version=metadata.get('name') + ' ' + metadata.get('version'), action='version')
	# misc
	group_misc = program.add_argument_group('misc')
	group_misc.add_argument('--skip-download', help=wording.get('help.skip_download'), action='store_true', default=config.get_bool_value('misc.skip_download'))
	group_misc.add_argument('--headless', help=wording.get('help.headless'), action='store_true', default=config.get_bool_value('misc.headless'))
	group_misc.add_argument('--log-level', help=wording.get('help.log_level'), default=config.get_str_value('misc.log_level', 'info'), choices=logger.get_log_levels())
	# execution
	execution_providers = encode_execution_providers(onnxruntime.get_available_providers())
	group_execution = program.add_argument_group('execution')
	group_execution.add_argument('--execution-providers', help=wording.get('help.execution_providers').format(choices=', '.join(execution_providers)), default=config.get_str_list('execution.execution_providers', 'cpu'), choices=execution_providers, nargs='+', metavar='EXECUTION_PROVIDERS')
	group_execution.add_argument('--execution-thread-count', help=wording.get('help.execution_thread_count'), type=int, default=config.get_int_value('execution.execution_thread_count', '4'), choices=facefusion.choices.execution_thread_count_range, metavar=create_metavar(facefusion.choices.execution_thread_count_range))
	group_execution.add_argument('--execution-queue-count', help=wording.get('help.execution_queue_count'), type=int, default=config.get_int_value('execution.execution_queue_count', '1'), choices=facefusion.choices.execution_queue_count_range, metavar=create_metavar(facefusion.choices.execution_queue_count_range))
	# memory
	group_memory = program.add_argument_group('memory')
	group_memory.add_argument('--video-memory-strategy', help=wording.get('help.video_memory_strategy'), default=config.get_str_value('memory.video_memory_strategy', 'strict'), choices=facefusion.choices.video_memory_strategies)
	group_memory.add_argument('--system-memory-limit', help=wording.get('help.system_memory_limit'), type=int, default=config.get_int_value('memory.system_memory_limit', '0'), choices=facefusion.choices.system_memory_limit_range, metavar=create_metavar(facefusion.choices.system_memory_limit_range))
	# face analyser
	group_face_analyser = program.add_argument_group('face analyser')
	group_face_analyser.add_argument('--face-analyser-order', help=wording.get('help.face_analyser_order'), default=config.get_str_value('face_analyser.face_analyser_order', 'left-right'), choices=facefusion.choices.face_analyser_orders)
	group_face_analyser.add_argument('--face-analyser-age', help=wording.get('help.face_analyser_age'), default=config.get_str_value('face_analyser.face_analyser_age'), choices=facefusion.choices.face_analyser_ages)
	group_face_analyser.add_argument('--face-analyser-gender', help=wording.get('help.face_analyser_gender'), default=config.get_str_value('face_analyser.face_analyser_gender'), choices=facefusion.choices.face_analyser_genders)
	group_face_analyser.add_argument('--face-detector-model', help=wording.get('help.face_detector_model'), default=config.get_str_value('face_analyser.face_detector_model', 'yoloface'), choices=facefusion.choices.face_detector_set.keys())
	group_face_analyser.add_argument('--face-detector-size', help=wording.get('help.face_detector_size'), default=config.get_str_value('face_analyser.face_detector_size', '640x640'))
	group_face_analyser.add_argument('--face-detector-score', help=wording.get('help.face_detector_score'), type=float, default=config.get_float_value('face_analyser.face_detector_score', '0.5'), choices=facefusion.choices.face_detector_score_range, metavar=create_metavar(facefusion.choices.face_detector_score_range))
	# face selector
	group_face_selector = program.add_argument_group('face selector')
	group_face_selector.add_argument('--face-selector-mode', help=wording.get('help.face_selector_mode'), default=config.get_str_value('face_selector.face_selector_mode', 'reference'), choices=facefusion.choices.face_selector_modes)
	group_face_selector.add_argument('--reference-face-position', help=wording.get('help.reference_face_position'), type=int, default=config.get_int_value('face_selector.reference_face_position', '0'))
	group_face_selector.add_argument('--reference-face-distance', help=wording.get('help.reference_face_distance'), type=float, default=config.get_float_value('face_selector.reference_face_distance', '0.6'), choices=facefusion.choices.reference_face_distance_range, metavar=create_metavar(facefusion.choices.reference_face_distance_range))
	group_face_selector.add_argument('--reference-frame-number', help=wording.get('help.reference_frame_number'), type=int, default=config.get_int_value('face_selector.reference_frame_number', '0'))
	# face mask
	group_face_mask = program.add_argument_group('face mask')
	group_face_mask.add_argument('--face-mask-types', help=wording.get('help.face_mask_types').format(choices=', '.join(facefusion.choices.face_mask_types)), default=config.get_str_list('face_mask.face_mask_types', 'box'), choices=facefusion.choices.face_mask_types, nargs='+', metavar='FACE_MASK_TYPES')
	group_face_mask.add_argument('--face-mask-blur', help=wording.get('help.face_mask_blur'), type=float, default=config.get_float_value('face_mask.face_mask_blur', '0.3'), choices=facefusion.choices.face_mask_blur_range, metavar=create_metavar(facefusion.choices.face_mask_blur_range))
	group_face_mask.add_argument('--face-mask-padding', help=wording.get('help.face_mask_padding'), type=int, default=config.get_int_list('face_mask.face_mask_padding', '0 0 0 0'), nargs='+')
	group_face_mask.add_argument('--face-mask-regions', help=wording.get('help.face_mask_regions').format(choices=', '.join(facefusion.choices.face_mask_regions)), default=config.get_str_list('face_mask.face_mask_regions', ' '.join(facefusion.choices.face_mask_regions)), choices=facefusion.choices.face_mask_regions, nargs='+', metavar='FACE_MASK_REGIONS')
	# frame extraction
	group_frame_extraction = program.add_argument_group('frame extraction')
	group_frame_extraction.add_argument('--trim-frame-start', help=wording.get('help.trim_frame_start'), type=int, default=facefusion.config.get_int_value('frame_extraction.trim_frame_start'))
	group_frame_extraction.add_argument('--trim-frame-end', help=wording.get('help.trim_frame_end'), type=int, default=facefusion.config.get_int_value('frame_extraction.trim_frame_end'))
	group_frame_extraction.add_argument('--temp-frame-format', help=wording.get('help.temp_frame_format'), default=config.get_str_value('frame_extraction.temp_frame_format', 'jpg'), choices=facefusion.choices.temp_frame_formats)
	group_frame_extraction.add_argument('--temp-frame-quality', help=wording.get('help.temp_frame_quality'), type=int, default=config.get_int_value('frame_extraction.temp_frame_quality', '100'), choices=facefusion.choices.temp_frame_quality_range, metavar=create_metavar(facefusion.choices.temp_frame_quality_range))
	group_frame_extraction.add_argument('--keep-temp', help=wording.get('help.keep_temp'), action='store_true', default=config.get_bool_value('frame_extraction.keep_temp'))
	# output creation
	group_output_creation = program.add_argument_group('output creation')
	group_output_creation.add_argument('--output-image-quality', help=wording.get('help.output_image_quality'), type=int, default=config.get_int_value('output_creation.output_image_quality', '80'), choices=facefusion.choices.output_image_quality_range, metavar=create_metavar(facefusion.choices.output_image_quality_range))
	group_output_creation.add_argument('--output-video-encoder', help=wording.get('help.output_video_encoder'), default=config.get_str_value('output_creation.output_video_encoder', 'libx264'), choices=facefusion.choices.output_video_encoders)
	group_output_creation.add_argument('--output-video-preset', help=wording.get('help.output_video_preset'), default=config.get_str_value('output_creation.output_video_preset', 'veryfast'), choices=facefusion.choices.output_video_presets)
	group_output_creation.add_argument('--output-video-quality', help=wording.get('help.output_video_quality'), type=int, default=config.get_int_value('output_creation.output_video_quality', '80'), choices=facefusion.choices.output_video_quality_range, metavar=create_metavar(facefusion.choices.output_video_quality_range))
	group_output_creation.add_argument('--output-video-resolution', help=wording.get('help.output_video_resolution'), default=config.get_str_value('output_creation.output_video_resolution'))
	group_output_creation.add_argument('--output-video-fps', help=wording.get('help.output_video_fps'), type=float)
	group_output_creation.add_argument('--skip-audio', help=wording.get('help.skip_audio'), action='store_true', default=config.get_bool_value('output_creation.skip_audio'))
	# frame processors
	available_frame_processors = list_directory('facefusion/processors/frame/modules')
	program = ArgumentParser(parents=[program], formatter_class=program.formatter_class, add_help=True)
	group_frame_processors = program.add_argument_group('frame processors')
	group_frame_processors.add_argument('--frame-processors', help=wording.get('help.frame_processors').format(choices=', '.join(available_frame_processors)), default=config.get_str_list('frame_processors.frame_processors', 'face_swapper'), nargs='+')
	for frame_processor in available_frame_processors:
		frame_processor_module = load_frame_processor_module(frame_processor)
		frame_processor_module.register_args(group_frame_processors)
	# uis
	available_ui_layouts = list_directory('facefusion/uis/layouts')
	group_uis = program.add_argument_group('uis')
	group_uis.add_argument('--ui-layouts', help=wording.get('help.ui_layouts').format(choices=', '.join(available_ui_layouts)), default=config.get_str_list('uis.ui_layouts', 'default'), nargs='+')

	apply_args(program)
	logger.init(facefusion.globals.log_level)
	if facefusion.globals.system_memory_limit > 0:
		limit_system_memory(facefusion.globals.system_memory_limit)
	if not pre_check() or not content_analyser.pre_check() or not face_analyser.pre_check() or not face_masker.pre_check():
		return
	for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors):
		if not frame_processor_module.pre_check():
			return


def download_resource(url):
    response = requests.get(url)
    response.raise_for_status()  # 确保请求成功

    _, ext = os.path.splitext(url)

    random_filename = f"{uuid.uuid4()}{ext}"

    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, random_filename)

    with open(temp_path, 'wb') as temp_file:
        temp_file.write(response.content)

    return temp_path


class ExceptionResponseSchema(BaseModel):
    error: str


class ChatRequestSchema(BaseModel):
    source_url: str = Field(..., description="source_url")
    target_url: str = Field(..., description="target_url")

    # 允许模型接受任何额外的字段
    class Config:
        extra = Extra.allow


face_router = APIRouter()


async def _face_swapper(args):
    try:
        source_url = args.get("source_url")
        target_url = args.get("target_url")
        random_filename = f"{uuid.uuid4()}.png"

        temp_dir = tempfile.gettempdir()
        temp_path = os.path.join(temp_dir, random_filename)

        facefusion.globals.source_paths = [download_resource(source_url)]
        facefusion.globals.target_path = download_resource(target_url)
        facefusion.globals.output_path = temp_path

        conditional_process()

        async with aiofiles.open(temp_path, 'rb') as f:
            image_data = await f.read()
        return image_data
        # return output_path
        # kwargs = {k: v for k, v in args.items() if k != "expr"}
        # return backtest(expr, **kwargs)
    except Exception as e:
        return e


@face_router.post("/swapper",
                  response_model=str,
                  responses={"400": {"model": ExceptionResponseSchema}})
async def swapper(request: ChatRequestSchema):
    request_id = str(uuid.uuid4())
    # logger.info(f'<{request_id}> [请求] {request}')
    # 将除了expr以外的所有字段作为kwargs传递
    kwargs = request.model_dump()

    res = await _face_swapper(kwargs)

    if isinstance(res, Exception):
        logger.error(f'[图片生成异常] {res}')
        raise HTTPException(status_code=400, detail=f'{res}')
    else:
        headers = {'Content-Disposition': 'attachment; filename="result.png"'}
        return Response(content=res, media_type='image/png', headers=headers)


run_model()